import numpy as np
import numba as nb

nu = 1.5e-5  # 运动粘性系数

@nb.njit(parallel=True)
def compute_velocity(X, Y, U, a):
    """计算速度场（Numba加速）"""
    Vx = np.zeros_like(X)
    Vy = np.zeros_like(Y)
    for i in nb.prange(X.shape[0]):
        for j in nb.prange(X.shape[1]):
            z = X[i, j] + 1j * Y[i, j]
            if np.abs(z) > a:
                V = U * (1 - (a**2) / (z**2))
                Vx[i, j] = V.real
                Vy[i, j] = V.imag
    return Vx, Vy

def calculate_stream_function(X, Y, U, a):
    """计算流函数ψ"""
    r_sq = X**2 + Y**2
    return U * (Y - (a**2 * Y) / np.where(r_sq != 0, r_sq, 1e-10))

from PyQt5.QtWidgets import QApplication, QMainWindow, QVBoxLayout, QWidget, QSlider, QLineEdit, QLabel
from PyQt5.QtCore import Qt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
import matplotlib.pyplot as plt

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.a = 1.0  # 初始半径
        self.U = 5.0  # 初始流速
        self.init_ui()
        self.update_plot()

    def init_ui(self):
        # 控件布局
        self.setWindowTitle("流场可视化系统")
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout(central_widget)

        # 圆柱半径滑动条
        self.slider_a = QSlider(Qt.Horizontal)
        self.slider_a.setRange(5, 20)  # 0.5~2.0（步长0.1）
        self.slider_a.setValue(10)
        self.slider_a.valueChanged.connect(self.update_a)
        layout.addWidget(QLabel("圆柱半径 (m):"))
        layout.addWidget(self.slider_a)

        # 流速输入框
        self.lineEdit_U = QLineEdit(str(self.U))
        self.lineEdit_U.editingFinished.connect(self.update_U)
        layout.addWidget(QLabel("流速 U (m/s):"))
        layout.addWidget(self.lineEdit_U)

        # 雷诺数显示
        self.label_Re = QLabel()
        layout.addWidget(self.label_Re)

        # 绘图区域
        self.figure = Figure()
        self.canvas = FigureCanvas(self.figure)
        layout.addWidget(self.canvas)

    def update_a(self):
        self.a = self.slider_a.value() / 10.0
        self.update_plot()

    def update_U(self):
        try:
            self.U = float(self.lineEdit_U.text())
        except ValueError:
            pass
        self.update_plot()

    def update_plot(self):
        a, U = self.a, self.U
        X, Y = np.meshgrid(np.linspace(-5, 5, 300), np.linspace(-5, 5, 300))
        Vx, Vy = compute_velocity(X, Y, U, a)
        psi = calculate_stream_function(X, Y, U, a)

        # 清空画布
        self.figure.clf()
        ax = self.figure.add_subplot(111)

        # 流线图（白色实线）
        ax.streamplot(X, Y, Vx, Vy, color='white', linewidth=0.8, density=1.5)

        # 等势线图（黄色虚线）
        ax.contour(X, Y, psi, levels=20, colors='yellow', linestyles='--', linewidths=0.8)

        # 圆柱填充（红色半透明）
        circle = plt.Circle((0, 0), a, color='red', alpha=0.3, zorder=3)
        ax.add_patch(circle)

        # 驻点标注（理论位置）
        ax.plot([a, -a], [0, 0], 'ro', markersize=4, zorder=4)

        # 图形设置
        ax.set_xlim(-5, 5)
        ax.set_ylim(-5, 5)
        ax.set_facecolor('k')
        ax.set_aspect('equal')

        # 雷诺数计算
        Re = 2 * U * a / nu
        self.label_Re.setText(f"雷诺数 Re = {Re:.1f}")

        self.canvas.draw()

if __name__ == '__main__':
    app = QApplication([])
    window = MainWindow()
    window.resize(800, 600)
    window.show()
    app.exec_()
    
    # 物理验证
def validate_streamlines(a, U):
    theta = np.linspace(0, 2*np.pi, 10)
    x = a * np.cos(theta)
    y = a * np.sin(theta)
    psi_values = calculate_stream_function(x, y, U, a)
    std = np.std(psi_values)
    return std < 0.01 * U / a
    